package com.autonavi.collapse.task.dispatch;

import com.autonavi.collapse.task.model.TaskConfig;
import com.autonavi.collapse.task.model.TaskRelation;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.stream.Collectors;

public class MinimumTaskDispatcher implements TaskDispatcher {

    private static final Logger LOGGER = LoggerFactory.getLogger(MinimumTaskDispatcher.class);


    @Override
    public List<TaskConfig> dispatch(DispatchContext dispatchContext) {
        if (dispatchContext == null) {
            return null;
        }
        List<TaskNode> availableNodes = dispatchContext.getAvailableNodes();
        if (availableNodes == null) {
            return null;
        }
        List<TaskConfig> tasks = new ArrayList<>();

        List<TaskConfig> deadTasks = dispatchContext.getDeadTasks();
        if (deadTasks != null) {
            tasks.addAll(deadTasks);
        }

        List<TaskConfig> newAndRetryTasks = dispatchContext.getTasks();
        if (newAndRetryTasks != null) {
            tasks.addAll(newAndRetryTasks);
        }

        List<TaskConfig> results = new ArrayList<>();

        List<TaskRelation> taskRelations = dispatchContext.getRelations();
        List<TaskConfig> fixedTasks = new ArrayList<>();
        List<TaskNode> fixedNodes = new ArrayList<>();
        if (taskRelations != null) {
            Map<String, List<TaskConfig>> taskMaps = tasks.stream().filter(e -> StringUtils.isNotBlank(e.getType())).collect(Collectors.groupingBy(TaskConfig::getType));
            if (taskMaps != null) {
                Set<String> tmpNodes = new HashSet<>();
                for (TaskRelation relation : taskRelations) {
                    Set<String> set = taskMaps.keySet();
                    if (set.contains(relation.getTaskType())) {
                        List<TaskConfig> list = taskMaps.get(relation.getTaskType());
                        if (list != null) {
                            fixedTasks.addAll(list);
                        }
                        tmpNodes.add(relation.getNodeName());
                    }
                }
                tasks.removeAll(fixedTasks);
                List<TaskNode> remainNodes = new ArrayList<>();
                if (!tmpNodes.isEmpty()) {
                    for (TaskNode node : availableNodes) {
                        if (tmpNodes.contains(node.getName())) {
                            fixedNodes.add(node);
                        } else {
                            remainNodes.add(node);
                        }
                    }
                    if (!remainNodes.isEmpty()) {
                        availableNodes = remainNodes;
                    }
                }
                if (!fixedNodes.isEmpty()) {
                    PriorityQueue<TaskNode> speciaNodes = new PriorityQueue<>((o1,o2) -> Integer.compare(o2.getFreeSize(), o1.getFreeSize()));
                    for (TaskNode fixedNode : fixedNodes) {
                        speciaNodes.offer(fixedNode);
                    }
                    List<TaskConfig> spList = doDispatch(speciaNodes, fixedTasks);
                    if (!spList.isEmpty()) {
                        results.addAll(spList);
                    }
                } else {
                    tasks.addAll(fixedTasks);
                }
            }
        }
        PriorityQueue<TaskNode> heapSortedNodes = new PriorityQueue<>((o1,o2) -> Integer.compare(o2.getFreeSize(), o1.getFreeSize()));
        for (TaskNode availableNode : availableNodes) heapSortedNodes.offer(availableNode);
        List<TaskConfig> dispatchTasks = doDispatch(heapSortedNodes, tasks);
        if (!dispatchTasks.isEmpty()) {
            results.addAll(dispatchTasks);
        }
        return results;
    }



    private List<TaskConfig> doDispatch(PriorityQueue<TaskNode> taskNodes, List<TaskConfig> tasks){
        List<TaskConfig> result = new ArrayList<>();
        for (TaskConfig task : tasks) {
            TaskNode taskNode = taskNodes.poll();
            if (taskNode != null && taskNode.getFreeSize() > 0) {
                task.setStatus(TaskConfig.DISPATCHED);
                task.setOwner(taskNode.getName());
                result.add(task);
                taskNode.setFreeSize(taskNode.getFreeSize() - 1);
                taskNodes.offer(taskNode);
            }
        }
        return result;
    }
}
